import pandas as pd
import numpy as np

def split_data_by_groups(file_path, group_size=5, save_data=True):
    """
    按照每group_size组取最后一组作为测试数据的规则分割数据
    
    参数:
    file_path: 数据文件路径
    group_size: 每组的大小，默认为5
    save_data: 是否保存分割后的数据，默认为True
    
    返回:
    X_train, y_train, X_test, y_test: 训练和测试数据
    """
    print(f"正在加载数据: {file_path}")
    
    # 读取Excel文件
    df = pd.read_excel(file_path)
    print(f"数据形状: {df.shape}")
    print(f"数据列名: {df.columns.tolist()}")
    
    # 提取电压序列和容量
    voltages_data = []
    capacities = []
    initial_capacities = []
    cycles = []
    rates = []
    temps = []
    
    for idx, row in df.iterrows():
        # 处理电压序列（字符串格式的列表）
        voltage_str = row['Voltages']
        if isinstance(voltage_str, str):
            # 将字符串转换为列表
            voltage_list = eval(voltage_str)
        else:
            voltage_list = voltage_str
            
        # 确保电压序列长度为14
        if len(voltage_list) == 14:
            voltages_data.append(voltage_list)
            capacities.append(row['Capacity'])
            initial_capacities.append(row['initial_capacity'])
            cycles.append(row['cycle'])
            rates.append(row['rate'])
            temps.append(row['Tem'])
            
    # 转换为numpy数组
    X = np.array(voltages_data, dtype=np.float32)
    y = np.array(capacities, dtype=np.float32)
    initial_capacities_arr = np.array(initial_capacities, dtype=np.float32)
    cycles_arr = np.array(cycles, dtype=np.int32)
    rates_arr = np.array(rates, dtype=np.float32)
    temps_arr = np.array(temps, dtype=np.float32)
    
    print(f"有效样本数: {len(X)}")
    
    # 按照每group_size组取最后一组作为测试数据的规则分割
    train_indices = []
    test_indices = []
    
    for i in range(0, len(X), group_size):
        group_end = min(i + group_size, len(X))
        
        # 训练数据：每组的前group_size-1个样本
        train_indices.extend(range(i, group_end - 1))
        
        # 测试数据：每组的最后一个样本
        if group_end - 1 >= i:  # 确保有最后一个样本
            test_indices.append(group_end - 1)
    
    X_train = X[train_indices]
    y_train = y[train_indices]
    initial_capacities_train = initial_capacities_arr[train_indices]
    cycles_train = cycles_arr[train_indices]
    rates_train = rates_arr[train_indices]
    temps_train = temps_arr[train_indices]
    
    X_test = X[test_indices]
    y_test = y[test_indices]
    initial_capacities_test = initial_capacities_arr[test_indices]
    cycles_test = cycles_arr[test_indices]
    rates_test = rates_arr[test_indices]
    temps_test = temps_arr[test_indices]
    
    print(f"训练样本数: {len(X_train)}")
    print(f"测试样本数: {len(X_test)}")
    print(f"分割比例: 训练 {len(X_train)/len(X)*100:.1f}%, 测试 {len(X_test)/len(X)*100:.1f}%")
    
    # 保存分割后的数据
    if save_data:
        # 创建训练数据DataFrame
        train_data = []
        for i in range(len(X_train)):
            train_data.append({
                'cycle': cycles_train[i],
                'Voltages': X_train[i].tolist(),
                'rate': rates_train[i],
                'Tem': temps_train[i],
                'Capacity': y_train[i],
                'initial_capacity': initial_capacities_train[i]
            })
        
        train_df = pd.DataFrame(train_data)
        train_file = "NCM_NCA_train_data.csv"
        train_df.to_csv(train_file, index=False)
        print(f"训练数据已保存到: {train_file}")
        
        # 创建测试数据DataFrame
        test_data = []
        for i in range(len(X_test)):
            test_data.append({
                'cycle': cycles_test[i],
                'Voltages': X_test[i].tolist(),
                'rate': rates_test[i],
                'Tem': temps_test[i],
                'Capacity': y_test[i],
                'initial_capacity': initial_capacities_test[i]
            })
        
        test_df = pd.DataFrame(test_data)
        test_file = "NCM_NCA_test_data.csv"
        test_df.to_csv(test_file, index=False)
        print(f"测试数据已保存到: {test_file}")
        

    
    return X_train, y_train, X_test, y_test

def load_train_test_data(file_path, group_size=5):
    """
    加载训练和测试数据的接口函数
    
    参数:
    file_path: 数据文件路径
    group_size: 每组的大小，默认为5(选取每组最后一个作为测试数据)
    
    返回:
    X_train, y_train, X_test, y_test: 训练和测试数据
    """
    return split_data_by_groups(file_path, group_size)

# 测试函数
if __name__ == "__main__":
    # 测试数据分割
    data_file = "data_set/Dataset_NCM_NCA_battery_with_initial_capacity.xlsx"
    
    try:
        X_train, y_train, X_test, y_test = load_train_test_data(data_file)
        
        print(f"\n训练数据形状: {X_train.shape}")
        print(f"测试数据形状: {X_test.shape}")
        print(f"训练容量范围: {y_train.min():.2f} - {y_train.max():.2f}")
        print(f"测试容量范围: {y_test.min():.2f} - {y_test.max():.2f}")
          
    except Exception as e:
        print(f"错误: {e}")
        print("请确保数据文件存在且格式正确")
